code: def compute_weight_mat(
      input_size: core.DimSize,
      output_size: core.DimSize,
      scale,
      translation,
      kernel: Callable,
      antialias: bool,
  ):
      dtype = jnp.result_type(scale, translation)
      inv_scale = 1.0 / scale
      # When downsampling the kernel should be scaled since we want to low pass
      # filter and interpolate, but when upsampling it should not be since we only
      # want to interpolate.
      kernel_scale = jnp.maximum(inv_scale, 1.0) if antialias else 1.0
      sample_f = (
          (jnp.arange(output_size, dtype=dtype) + 0.5) * inv_scale
          - translation * inv_scale
          - 0.5
      )
      x = (
          jnp.abs(
              sample_f[jnp.newaxis, :]
              - jnp.arange(input_size, dtype=dtype)[:, jnp.newaxis]
          )
          / kernel_scale
      )
      weights = kernel(x)
  
      total_weight_sum = jnp.sum(weights, axis=0, keepdims=True)
      weights = jnp.where(
          jnp.abs(total_weight_sum)
          > 1000.0 * float(np.finfo(np.float32).eps),
          jnp.divide(
              weights,
              jnp.where(total_weight_sum != 0, total_weight_sum, 1),
          ),
          0,
      )
      # Zero out weights where the sample location is completely outside the input
      # range.
      # Note sample_f has already had the 0.5 removed, hence the weird range below.
      input_size_minus_0_5 = core.dimension_as_value(input_size) - 0.5
      return jnp.where(
          jnp.logical_and(
              sample_f >= -0.5, sample_f <= input_size_minus_0_5
          )[jnp.newaxis, :],
          weights,
          0,
      )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  from functools import partial
  import enum
  from typing import Callable
  import numpy
  from jax import jit
  from jax import lax
  from jax import numpy
  from jax._src import core
  from jax._src.util import canonicalize_axis
  from jax._src.numpy.util import promote_dtypes_inexact
label_desc: Use image node_id 4 for calling a global function with example usage: compute_weight_mat(input_size, output_size, scale, translation, kernel, antialias) and returns: jnp