Fix torchax tensor flatten bug on empty tensors

[cychiuak]

Flatten on tensor of shape contains zero is supported in PyTorch, but not through torchax.

import torch
import torchax

with torchax.default_env():
    value = torch.ones((16, 0))
    print(value.flatten(0, -2).shape) # works in CPU > torch.Size([16, 0])
    value = value.to(device='jax')
    print(value.flatten(0, -2)) # throws exception

This is due to jax.numpy fails to calculate the correct dimension when input dimension contains 0 and -1 at the same time. It returns a division by zero error.

import jax.numpy as jnp

print(jnp.reshape(jnp.ones((2, 3)), (-1, 3)).shape) # works > (2, 3)
value = jnp.ones((16, 0))
print(jnp.reshape(value, (16, 0)).shape) # works > (16, 0)
print(jnp.reshape(value, (-1, 0)).shape) # should work but throws exception

The new fix avoids sending -1 into jnp.shape by explicitly calculating the flattening dimension to to fix this error:

flattened_size = math.prod(self._elem.shape[start_dim : end_dim + 1])
new_shape = self._elem.shape[:start_dim] + (flattened_size,) + self._elem.shape[end_dim + 1 :]

[wdhongtw]

LGTM! 👍

Just a few suggestions.