Repository: keras-team/keras

Issue:
ops.take results in different tensor shape in tensorflow and torch (jax is the same as tensorflow)
I come across a strange problem when using `ops.take` in different backends:

```python
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

from keras import ops

num = ops.reshape(ops.arange(10), (2, 5))
print(ops.take(num, ops.where(ops.greater(num, 5))))
```
```
tf.Tensor(
[[1 1 1 1]
 [1 2 3 4]], shape=(2, 4), dtype=int32)
```

Change the backend to `jax`, it works the same as `tensorflow`:
```
Array([[1, 1, 1, 1],
       [1, 2, 3, 4]], dtype=int32)
```

Change the backend to `torch` the print shows:
```
../aten/src/ATen/native/cuda/Indexing.cu:1239: indexSelectSmallIndex: block: [0,0,0], thread: [0,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1239: indexSelectSmallIndex: block: [0,0,0], thread: [1,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1239: indexSelectSmallIndex: block: [0,0,0], thread: [2,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1239: indexSelectSmallIndex: block: [0,0,0], thread: [3,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
../aten/src/ATen/native/cuda/Indexing.cu:1239: indexSelectSmallIndex: block: [0,0,0], thread: [4,0,0] Assertion `srcIndex < srcSelectDimSize` failed.

{
	"name": "RuntimeError",
	"message": "CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
",
	"stack": "---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File ~/miniconda3/lib/python3.10/site-packages/IPython/core/formatters.py:708, in PlainTextFormatter.__call__(self, obj)
    701 stream = StringIO()
    702 printer = pretty.RepresentationPrinter(stream, self.verbose,
    703     self.max_width, self.newline,
    704     max_seq_length=self.max_seq_length,
    705     singleton_pprinters=self.singleton_printers,
    706     type_pprinters=self.type_printers,
    707     deferred_pprinters=self.deferred_printers)
--> 708 printer.pretty(obj)
    709 printer.flush()
    710 return stream.getvalue()

File ~/miniconda3/lib/python3.10/site-packages/IPython/lib/pretty.py:410, in RepresentationPrinter.pretty(self, obj)
    407                         return meth(obj, self, cycle)
    408                 if cls is not object \\
    409                         and callable(cls.__dict__.get('__repr__')):
--> 410                     return _repr_pprint(obj, self, cycle)
    412     return _default_pprint(obj, self, cycle)
    413 finally:

File ~/miniconda3/lib/python3.10/site-packages/IPython/lib/pretty.py:778, in _repr_pprint(obj, p, cycle)
    776 \"\"\"A pprint that just redirects to the normal repr function.\"\"\"
    777 # Find newlines and replace them with p.break_()
--> 778 output = repr(obj)
    779 lines = output.splitlines()
    780 with p.group():

File ~/miniconda3/lib/python3.10/site-packages/torch/_tensor.py:431, in Tensor.__repr__(self, tensor_contents)
    427     return handle_torch_function(
    428         Tensor.__repr__, (self,), self, tensor_contents=tensor_contents
    429     )
    430 # All strings are unicode in Python 3.
--> 431 return torch._tensor_str._str(self, tensor_contents=tensor_contents)

File ~/miniconda3/lib/python3.10/site-packages/torch/_tensor_str.py:664, in _str(self, tensor_contents)
    662 with torch.no_grad(), torch.utils._python_dispatch._disable_current_modes():
    663     guard = torch._C._DisableFuncTorch()
--> 664     return _str_intern(self, tensor_contents=tensor_contents)

File ~/miniconda3/lib/python3.10/site-packages/torch/_tensor_str.py:595, in _str_intern(inp, tensor_contents)
    593                     tensor_str = _tensor_str(self.to_dense(), indent)
    594                 else:
--> 595                     tensor_str = _tensor_str(self, indent)
    597 if self.layout != torch.strided:
    598     suffixes.append(\"layout=\" + str(self.layout))

File ~/miniconda3/lib/python3.10/site-packages/torch/_tensor_str.py:347, in _tensor_str(self, indent)
    343     return _tensor_str_with_formatter(
    344         self, indent, summarize, real_formatter, imag_formatter
    345     )
    346 else:
--> 347     formatter = _Formatter(get_summarized_data(self) if summarize else self)
    348     return _tensor_str_with_formatter(self, indent, summarize, formatter)

File ~/miniconda3/lib/python3.10/site-packages/torch/_tensor_str.py:133, in _Formatter.__init__(self, tensor)
    131 if not self.floating_dtype:
    132     for value in tensor_view:
--> 133         value_str = f\"{value}\"
    134         self.max_width = max(self.max_width, len(value_str))
    136 else:

File ~/miniconda3/lib/python3.10/site-packages/torch/_tensor.py:933, in Tensor.__format__(self, format_spec)
    931     return handle_torch_function(Tensor.__format__, (self,), self, format_spec)
    932 if self.dim() == 0 and not self.is_meta and type(self) is Tensor:
--> 933     return self.item().__format__(format_spec)
    934 return object.__format__(self, format_spec)

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
"
}
```

Then I go for another test:
```python
num = ops.arange(3)
i, j = ops.meshgrid(num, num)
mask = ops.where(ops.greater(i, j))
print(ops.take(i, mask))
```

This time `torch` works fine but the results are different from those of `tensorflow` and `jax`:
```
# tensorflow
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[0, 0, 1],
       [1, 2, 2]], dtype=int32)>

# jax
Array([[0, 0, 1],
       [1, 2, 2]], dtype=int32)

# torch
tensor([[[0, 1, 2],
         [0, 1, 2],
         [0, 1, 2]],

        [[0, 1, 2],
         [0, 1, 2],
         [0, 1, 2]]], device='cuda:0', dtype=torch.int32)
```

Please check them. I want to implement some calculation in a layer's `call` for all backends. Thank you so much!


Discussion:
Hi dear keras team, I check the online doc about `ops.take` carefully and do a little more investigation. Here's my conclusion:

`ops.take` in `torch` backend doesn't flatten the inputs by default as the doc (https://keras.io/api/ops/numpy/#take-function) says.

Below is my colab script to produce it:
```python
import os
os.environ["KERAS_BACKEND"] = "jax"

from keras import ops

num = ops.reshape(ops.arange(10), (2, 5))

print(ops.take(num, 0))
print(ops.take(num, 6))
```

```
# tensorflow or jax
tf.Tensor(0, shape=(), dtype=int32)

# torch
tensor([0, 1, 2, 3, 4], dtype=torch.int32)
```

Tensorflow and Jax will flatten the inputs first, so the index 0 will give the exactly the first number, but Torch will generate the first row. The index 6 is valid for flat inputs but not the case for the original one.


This should be the bug actually.

Resolving patch:
diff --git a/keras/backend/torch/numpy.py b/keras/backend/torch/numpy.py
--- a/keras/backend/torch/numpy.py
+++ b/keras/backend/torch/numpy.py
@@ -1217,9 +1217,12 @@ def swapaxes(x, axis1, axis2):
 def take(x, indices, axis=None):
     x = convert_to_tensor(x)
     indices = convert_to_tensor(indices).long()
-    if x.ndim == 2 and (axis is None or axis == 0):
+    if x.ndim == 2 and axis == 0:
         # This case is equivalent to embedding lookup.
         return torch.nn.functional.embedding(indices, x)
+    if axis is None:
+        x = torch.reshape(x, (-1,))
+        axis = 0
     if axis is not None:
         # make sure axis is non-negative
         axis = len(x.shape) + axis if axis < 0 else axis


Test patch:
diff --git a/keras/ops/numpy_test.py b/keras/ops/numpy_test.py
--- a/keras/ops/numpy_test.py
+++ b/keras/ops/numpy_test.py
@@ -2476,6 +2476,12 @@ def test_take(self):
             knp.take(x, indices, axis=-2),
             np.take(x, indices, axis=-2),
         )
+        # test with axis=None & x.ndim=2
+        x = np.array(([1, 2], [3, 4]))
+        indices = np.array([2, 3])
+        self.assertAllClose(
+            knp.take(x, indices, axis=None), np.take(x, indices, axis=None)
+        )
 
     @parameterized.named_parameters(
         named_product(