code: def _as_matrix(quat: jax.Array) -> jax.Array:
      x = quat[0]
      y = quat[1]
      z = quat[2]
      w = quat[3]
      x2 = x * x
      y2 = y * y
      z2 = z * z
      w2 = w * w
      xy = x * y
      zw = z * w
      xz = x * z
      yw = y * w
      yz = y * z
      xw = x * w
      return jnp.array(
          [
              [+x2 - y2 - z2 + w2, 2 * (xy - zw), 2 * (xz + yw)],
              [2 * (xy + zw), -x2 + y2 - z2 + w2, 2 * (yz - xw)],
              [2 * (xz - yw), 2 * (yz + xw), -x2 - y2 + z2 + w2],
          ]
      )
  
imports: from __future__ import annotations
  import functools
  import re
  import typing
  import scipy.spatial.transform
  import jax
  import jax.numpy
  from jax._src.numpy.util import _wraps
label_desc: Use image node_id 2 for calling a global function with example usage: _as_matrix(quat) and returns: jnp