From d6f390cc7b721c40d4881717e4567cf52d5f9a87 Mon Sep 17 00:00:00 2001 From: linmin Date: Tue, 2 Sep 2025 14:40:12 +0800 Subject: [PATCH] alphabetical order --- torch2jax/__init__.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/torch2jax/__init__.py b/torch2jax/__init__.py index f84a980..2b8cd5b 100644 --- a/torch2jax/__init__.py +++ b/torch2jax/__init__.py @@ -400,6 +400,12 @@ def flatten(input, start_dim=0, end_dim=-1): return jnp.reshape(_v(input), input.shape[:start_dim] + (-1,)) +@implements(torch.mean, Torchish_member=True, out_kwarg=True) +def mean(input, dim=None, keepdim=False, dtype=None): + dtype = t2j_dtype(dtype) if dtype is not None else None + return jnp.mean(_v(input), axis=dim, keepdims=keepdim, dtype=dtype) + + @implements(torch.multinomial, out_kwarg=True, Torchish_member=True) def multinomial(input, num_samples, replacement=False, generator=None): assert generator is None, "TODO: implement `generator`" @@ -418,12 +424,6 @@ def multinomial(input, num_samples, replacement=False, generator=None): raise ValueError(f"unsupported shape: {input.shape}") -@implements(torch.mean, Torchish_member=True, out_kwarg=True) -def mean(input, dim=None, keepdim=False, dtype=None): - dtype = t2j_dtype(dtype) if dtype is not None else None - return jnp.mean(_v(input), axis=dim, keepdims=keepdim, dtype=dtype) - - @implements(torch.normal, out_kwarg=True) def normal(*args, **kwargs): assert kwargs.get("generator", None) is None, "TODO: implement `generator`"