code: def _csr_transpose(data, indices, indptr):
      # Transpose of a square CSR matrix
      m = indptr.size - 1
      row = jnp.cumsum(jnp.zeros_like(indices).at[indptr].add(1)) - 1
      row_T, indices_T, data_T = jax.lax.sort(
          (indices, row, data), num_keys=2
      )
      indptr_T = (
          jnp.zeros_like(indptr)
          .at[1:]
          .set(
              jnp.cumsum(jnp.bincount(row_T, length=m)).astype(
                  indptr.dtype
              )
          )
      )
      return data_T, indices_T, indptr_T
  
imports: from __future__ import annotations
  from typing import Union, Callable
  import functools
  import jax
  import jax.numpy
  from jax.experimental import sparse
  from jax.interpreters import mlir
  from jax.interpreters import xla
  from jax._src import core
  from jax._src.interpreters import ad
  from jax._src.lib import gpu_solver
  import numpy
  from scipy.sparse import csr_matrix, linalg
label_desc: Use image node_id 18 for calling a global function with example usage: _csr_transpose(data, indices, indptr) and returns: data_T,  indices_T,  indptr_T