Skip to content

Commit

Permalink
Return a tuple from ops.shape with the Torch backend.
Browse files Browse the repository at this point in the history
With Torch, `x.shape` returns a `torch.Size`, which is a subclass of `tuple` but can cause different behaviors. In particular `convert_to_tensor` does not work on `torch.Size`.

This fixes keras-team#18900
  • Loading branch information
hertschuh committed Apr 29, 2024
1 parent 4cb5671 commit 4cc6064
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 2 deletions.
3 changes: 2 additions & 1 deletion keras/src/backend/torch/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -234,7 +234,8 @@ def is_tensor(x):


def shape(x):
return x.shape
# Convert from `torch.Size` to plain tuple.
return tuple(x.shape)


def cast(x, dtype):
Expand Down
4 changes: 3 additions & 1 deletion keras/src/ops/core_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,10 +338,12 @@ def test_stop_gradient_return(self):
self.assertAllClose(x, y)

def test_shape(self):
x = np.ones((2, 3, 7, 1))
x = ops.ones((2, 3, 7, 1))
self.assertEqual(core.shape(x).__class__, tuple)
self.assertAllEqual(core.shape(x), (2, 3, 7, 1))

x = KerasTensor((None, 3, None, 1))
self.assertEqual(core.shape(x).__class__, tuple)
self.assertAllEqual(core.shape(x), (None, 3, None, 1))

@pytest.mark.skipif(
Expand Down

0 comments on commit 4cc6064

Please sign in to comment.