code: def _lu(
      input, output_idx_type=np.int32, name=None
  ):  # pylint: disable=redefined-builtin
      """Returns Lu(lu, p), as TF does."""
      del name
      input = ops.convert_to_tensor(input)
      if (
          JAX_MODE
      ):  # JAX uses XLA, which can do a batched factorization.
          lu_out, pivots = scipy_linalg.lu_factor(input)
          from jax import lax  # pylint: disable=g-import-not-at-top
  
          return Lu(
              lu_out,
              lax.linalg.lu_pivots_to_permutation(
                  pivots, lu_out.shape[-1]
              ),
          )
      # Scipy can't batch, so we must do so manually.
      nbatch = int(np.prod(input.shape[:-2]))
      dim = input.shape[-1]
      flat_mat = input.reshape(nbatch, dim, dim)
      flat_lu = np.empty((nbatch, dim, dim), dtype=input.dtype)
      flat_piv = np.empty(
          (nbatch, dim), dtype=utils.numpy_dtype(output_idx_type)
      )
      if np.size(flat_lu):  # Avoid non-empty batches of empty matrices.
          for i, mat in enumerate(flat_mat):
              lu_out, pivots = scipy_linalg.lu_factor(mat)
              flat_lu[i] = lu_out
              flat_piv[i] = _lu_pivot_to_permutation(
                  pivots, flat_lu.shape[-1]
              )
      return Lu(
          flat_lu.reshape(*input.shape),
          flat_piv.reshape(*input.shape[:-1]),
      )
  
imports: import collections
  import functools
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
label_desc: Use image node_id 11 for calling a global function with example usage: _lu(input, output_idx_type, name) and returns: Lu, Lu