Repository: pytorch
Issue: Incompatibility with Numpy,JAX and MLX for int and bool type inputs
Labels: triaged,module: numpy,module: python frontend

### 🚀 The feature, motivation and pitch

There are some functions in Torch which behaves differently or raises errors where Numpy, JAX and MLX works with no differences or errors. I provided a script which depicts these issues in Additional context.  In this section one case may summarize the issue.

In Torch one can easily get average of two integer type tensor as:
```
print((torch.tensor([1]) + torch.tensor([2])) / 2)
#Output: tensor([1.5000])
```
but when same operation is performed for one tensor with two integer entries
```
print(torch.tensor([1, 2]).mean())
#Output: RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long
```
raises RuntimeError which is not the case for Numpy, JAX or MLX. 

Using provided script you can test several cases that I observed. For example, AvgPool1d casts result to int while JAX or MLX returns float results which is meaningful for averaging operation.

### Alternatives

It would be nice to get same results for commonly used frameworks like Numpy, JAX and MLX.

### Additional context

```
import flax.linen
import numpy as np
import jax.numpy as jnp
import flax
import mlx.nn.layers as mlx_nn
import mlx.core as mx
import torch

# Create a reference bool type tensor in numpy.
np_tensor = np.array(
    [[True, False, True], 
     [False, False, True]
    ]
)
# # int type reference.
# np_tensor = np.array([[1, 0, 1], [0, 0, 1]])

# Create same tensor in all frameworks.
jax_tensor = jnp.array(np_tensor)
torch_tensor = torch.tensor(np_tensor)
mlx_tensor = mx.array(np_tensor)

# Print sum operation for default axis. No differences here.
print(np_tensor.sum()) # Output: 3
print(jax_tensor.sum()) # Output: 3
print(mlx_tensor.sum()) # Output: array(3, dtype=int64)
print(torch_tensor.sum()) # Output: tensor(3)

# Print mean operation for default axis. Torch raises RuntimeError.
print(np_tensor.mean()) # Output: 3
print(jax_tensor.mean()) # Output: 3
print(mlx_tensor.mean()) # Output: array(0.5, dtype=float32)
# print(torch_tensor.mean()) 
# Output: RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Long


# Another difference can be observed for AvgPool1d operation.
input_tensor = np.array([[[1, 4], [2, 5], [3, 6]]])

# Flax output.
flax_out = flax.linen.avg_pool(input_tensor, (2,))
print(flax_out)
# Output: 
# [[[1.5 4.5]
#   [2.5 5.5]]]


# MLX output.
mlx_out = mlx_nn.AvgPool1d(kernel_size = 2, stride = 1)(mx.array(input_tensor))
print(mlx_out) 
# Output:
# array([[[1.5, 4.5],
#         [2.5, 5.5]]], dtype=float32)


# Torch output.
torch_out = torch.nn.AvgPool1d(kernel_size = (2,), stride = 1)(torch.tensor(input_tensor.transpose(0, 2, 1)))
print(torch_out.permute(0, 2, 1)) 
# Output: casts all results to int!
# tensor([[[1, 4],
#          [2, 5]]])

# # Torch also raises RuntimeError for MaxPool1d for int type input.
# torch_out = torch.nn.MaxPool1d(kernel_size = (2,), stride = 1)(torch.tensor(input_tensor.transpose(0, 2, 1)))
# RuntimeError: "max_pool1d_impl" not implemented for 'Long'


# Also LeakyReLU and gelu are different too.
input_tensor = np.array([[-1, 0], [3, -3]])

# Flax output.
flax_out = flax.linen.leaky_relu(input_tensor)
print(flax_out)
# Output: 
# [[-0.01  0.  ]
#  [ 3.   -0.03]]
flax_out = flax.linen.gelu(input_tensor, approximate = False)
print(flax_out)
# Output: 
# [[-0.15865526  0.        ]
#  [ 2.9959502  -0.00404969]]



# MLX output.
mlx_out = mlx_nn.LeakyReLU()(mx.array(input_tensor))
print(mlx_out) 
# Output:
# array([[-0.01, 0],
#        [3, -0.03]], dtype=float32)
mlx_out = mlx_nn.GELU()(mx.array(input_tensor))
print(mlx_out) 
# Output: 
# array([[-0.158655, 0],
#        [2.99595, -0.00404969]], dtype=float32)



# # Torch output.
# torch_out = torch.nn.LeakyReLU()(torch.tensor(input_tensor))
torch_out = torch.nn.GELU()(torch.tensor(input_tensor))
# RuntimeError: "leaky_relu_cpu" not implemented for 'Long'

```

cc @mruberry @rgommers @albanD