code: def _reshape_solve_result(result, b_batch_dims, result_permutation):
      """Reshapes result of a solve to the original shapes."""
      import jax.numpy as jnp  # pylint: disable=g-import-not-at-top
      import numpy as onp  # pylint: disable=g-import-not-at-top,reimported
  
      if b_batch_dims is None:
          return result
  
      new_shape = list(result.shape[:-1])
      new_shape.append(
          result.shape[-1] // int(onp.prod(onp.array(b_batch_dims)))
      )
      new_shape.extend(b_batch_dims)
      new_result = result.reshape(*new_shape)
      return jnp.transpose(new_result, result_permutation)
  
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 3 for calling a global function with example usage: _reshape_solve_result(result, b_batch_dims, result_permutation) and returns: jnp, result