code: def sytrd_hlo(dtype, a, *, lower):
      _lapack.initialize()
      a_type = ir.RankedTensorType(a.type)
      dims = a_type.shape
      assert len(dims) >= 2
      m, n = dims[-2:]
      assert m == n, (m, n)
      batch_dims = tuple(dims[:-2])
      num_bd = len(batch_dims)
      b = 1
      for d in batch_dims:
          b *= d
  
      if dtype == np.float32:
          fn = "lapack_ssytrd"
          lwork = _lapack.lapack_ssytrd_workspace(n, n)
          diag_type = a_type.element_type
      elif dtype == np.float64:
          fn = "lapack_dsytrd"
          lwork = _lapack.lapack_dsytrd_workspace(n, n)
          diag_type = a_type.element_type
      elif dtype == np.complex64:
          fn = "lapack_chetrd"
          lwork = _lapack.lapack_chetrd_workspace(n, n)
          diag_type = ir.F32Type.get()
      elif dtype == np.complex128:
          fn = "lapack_zhetrd"
          lwork = _lapack.lapack_zhetrd_workspace(n, n)
          diag_type = ir.F64Type.get()
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
      i32_type = ir.IntegerType.get_signless(32)
      out = custom_call(
          fn,
          result_types=[
              a.type,
              ir.RankedTensorType.get(batch_dims + (n,), diag_type),
              ir.RankedTensorType.get(batch_dims + (n - 1,), diag_type),
              ir.RankedTensorType.get(
                  batch_dims + (n - 1,), a_type.element_type
              ),
              ir.RankedTensorType.get(batch_dims, i32_type),
              ir.RankedTensorType.get([lwork], a_type.element_type),
          ],
          operands=[
              hlo_s32(n),
              hlo_s32(1 if lower else 0),
              hlo_s32(max(1, n)),
              hlo_s32(b),
              hlo_s32(lwork),
              a,
          ],
          operand_layouts=[[]] * 5 + [layout],
          result_layouts=[
              layout,
              (num_bd,) + tuple(range(num_bd - 1, -1, -1)),
              (num_bd,) + tuple(range(num_bd - 1, -1, -1)),
              (num_bd,) + tuple(range(num_bd - 1, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
              [0],
          ],
          operand_output_aliases={5: 0},
      ).results
      return out[:5]
  
file_path: jax/jaxlib/lapack.py
imports: from collections.abc import Sequence
  import numpy
  import jaxlib.mlir.ir
  import jaxlib.mlir.dialects.stablehlo
  from jaxlib import xla_client
  from .hlo_helpers import custom_call, hlo_u8, hlo_s32, ensure_hlo_s32, hlo_add, hlo_min, DimensionSize, ShapeTypePair, mk_result_types_and_shapes
  from .cpu import _lapack
label_desc: Use image node_id 11 for calling a global function with example usage: sytrd_hlo(dtype, a) and returns: out