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 5d610a2
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 3 deletions.
6 changes: 4 additions & 2 deletions 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 Expand Up @@ -290,7 +291,8 @@ def symbolic_call(fn, args, kwargs, fill_value):
(args, kwargs),
)
return fn(*meta_args, **meta_kwargs)
except:
except Exception as e:
print("### e", e)
with device_scope(DEFAULT_DEVICE):
# If the `"meta"` device placement fails, fall back to tracing
# eagerly with tensors on the default device. This will be
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 5d610a2

Please sign in to comment.