Generalize (join|split)_dims to work with arbitrary axes locations

Description

As discussed in https://github.com/pymc-devs/pytensor/pull/1842#issuecomment-3740698774

Instead of forcing the input axis to be consecutive in join_dims, and the output axis to come out consecutive in split_dims, we could generalize it so they can map anywhere and still be functional inverses of each other.

def join_dims(x, axes: int | Sequence[int] | None = None, output_axis: int = 0):
  if axes = None:
    axes = tuple(range(x.ndim))
  if isinstance(axes, int):
    axes = (axes,)  # No dims actually joined

  # I am not sure I want to do that, but if we could make output_axis default to the first of `axes` if these are consecutive
  # or zero otherwise. Much like numpy advanced indexing decides where to place the advanced views of the array.
  # This would be a more back-compatible with the old behavior of `join_dims` that only supported consecutive axis
  
  return JoinDims(output_axis, len(axes))(pt.moveaxis(x, axes, output_axis))


def split_dims(x, shape, axis: int = 0, output_axes: int | Sequence[int] | None = None):
  if output_axes = None:
    # default is the same location as input axis
    output_axes = axis
  if isinstance(output_axes, int):
    output_axes = tuple(range(output_axes, output_axes + len(shape))
  
  axes = tuple(range(axis, axis + len(shape))
  return pt.moveaxis(SplitDims(axis)(x), axes, output_axes)

We should decide pretty soon, as it's a breaking change, and we want to bring these ops to the spotlight.

Start with the join_dims and split_dims entry points and read the discussion in PR 1842, especially the linked comment. Determine the intended axis-mapping API and compatibility behavior before changing these operations. Done means arbitrary input and output axis locations are supported and join_dims and split_dims remain functional inverses.