aten::searchsorted lowering drops every keyword argument (right, side, out_int32, sorter)

## Summary

`torch.searchsorted(sorted_sequence, values, right=True)` raises `TypeError` under torchax. The
op **is** registered, so it does not fail with `OperatorNotFound`; the registered function
accepts only two positional parameters, while the aten schema has four keyword arguments.

`torchax/ops/jaten.py:277-286` (v0.0.13, unchanged on `main`):

```python
@op(torch.ops.aten.searchsorted.Tensor)
def _aten_searchsorted(sorted_sequence, values):
mappings.t2j_dtype(torch.get_default_dtype())
res = jnp.searchsorted(sorted_sequence, values)
if sorted_sequence.dtype == np.dtype(np.int32) or sorted_sequence.dtype == np.dtype(
np.int32
):
# res = res.astype(new_dtype)
res = res.astype(np.dtype(np.int64))
return res # jnp.searchsorted(sorted_sequence, values)
```

The aten schema is:

```
aten::searchsorted.Tensor(Tensor sorted_sequence, Tensor self, *, bool out_int32=False,
bool right=False, str? side=None, Tensor? sorter=None) -> Tensor
```

`jnp.searchsorted` supports the important one directly — `side='left'|'right'` — so this is a
signature gap, not a capability gap.

Three separate problems in those ten lines:

1. **`right` / `side` are unsupported**, so a caller asking for right-open bucketing silently
cannot express it, and gets a `TypeError` rather than a clear "unsupported argument".
2. **`out_int32` and `sorter` are likewise dropped.**
3. **The result dtype is forced to `int64` when the *haystack* is int32** (`:281-285`), which is
the opposite of `out_int32`'s meaning and, with `jax_enable_x64=False` (the default — see
below), produces a dtype JAX will not materialise at full width anyway. Line `:279` is a
no-op whose return value is discarded, and `:284` is commented-out code — this lowering looks
unfinished rather than deliberate.

Consistent with that reading, torchax's own conformance suite **skips** the op:
`test/test_ops.py:52` lists `"searchsorted"` in the skiplist, so the mismatch is known and
excluded rather than fixed.

## How it manifests

`torchax/tensor.py` dispatch calls `op.func(*args, **kwargs)` with the caller's kwargs passed
straight through and no per-op signature filtering, so a registered-but-narrower op raises a
plain Python `TypeError: _aten_searchsorted() got an unexpected keyword argument 'right'`. That
is harder to triage than a missing registration, because `OperatorNotFound` names the problem
and a `TypeError` points at torchax internals.

## Why it matters downstream

Encountered while running a PyTorch model on TPU through `tpu-inference` (which pins
`torchax==0.0.13`): the model used `torch.searchsorted(..., right=True)` for bucket lookup and
had to be rewritten to avoid it. Any model doing bucketing, histogram, or quantile-style index
lookup hits this, and the rewrite is not always semantics-preserving at bucket boundaries —
`right=True` is precisely the boundary behaviour.

## Suggested fix

```python
@op(torch.ops.aten.searchsorted.Tensor)
def _aten_searchsorted(sorted_sequence, values, *, out_int32=False, right=False,
side=None, sorter=None):
if sorter is not None:
sorted_sequence = jnp.take_along_axis(sorted_sequence, sorter, axis=-1)
if side is None:
side = "right" if right else "left"
elif right and side != ("right" if right else "left"):
raise ValueError("torch.searchsorted: `side` and `right` disagree")
res = jnp.searchsorted(sorted_sequence, values, side=side)
return res.astype(jnp.int32 if out_int32 else jnp.int64)
```

…and remove `"searchsorted"` from the `test/test_ops.py` skiplist so the OpInfo conformance test
covers it. Happy to send this as a PR with the skiplist entry removed if the shape looks right.

## Separately: the int64 half of what we originally reported

We had recorded "torchax has no int64". That is **not accurate** and we are correcting it here
rather than filing it:

- `torchax/ops/mappings.py:101-102` maps `torch.int64 -> jnp.int64` and `torch.long -> jnp.int64`
one-to-one. torchax does not downcast.
- 64-bit width is governed by JAX's global `jax_enable_x64`, which torchax exposes but leaves
off by default: `torchax/__init__.py:145` (`enable_accuracy_mode()` sets it `True`) versus
`:151` (`enable_performance_mode()` sets it `False`).

So the truncation we observed (`Explicitly requested dtype int64 ... truncated to int32`) comes
from JAX's x64 default, and — per JAX's own documentation on TPU — 64-bit integer support on
XLA:TPU is limited independently of that flag. A torchax-level fix would not restore it. Nothing
to file against torchax; the honest statement for a downstream backend is "assume 32-bit
integers on TPU and carry 64-bit arithmetic in limbs", which is what we did.

Start in torchax/ops/jaten.py:277-286 and compare the registered searchsorted lowering with the aten schema and JAX searchsorted arguments. Check torchax/tensor.py dispatch behavior, then update test/test_ops.py to remove searchsorted from the skiplist and run its conformance coverage. Done means right/side, out_int32, and sorter are handled consistently, result dtype follows out_int32, and the relevant tests pass.

Written by the indexing model from the issue text.