code: def convert_conv_or_linear_layer(keras_layer, objax_layer):
      """Converts variables of convolution/linear layer from Keras to Objax."""
      objax_layer.w = objax.TrainVar(
          jn.array(keras_layer.__dict__["kernel"].numpy())
      )
      if keras_layer.__dict__.get("bias") is not None:
          if jn.ndim(objax_layer.w.value) == 4:
              # conv layer
              bias_shape = (objax_layer.w.value.shape[-1], 1, 1)
          else:
              # linear layer
              bias_shape = objax_layer.w.value.shape[-1]
          objax_layer.b = objax.TrainVar(
              jn.array(keras_layer.__dict__["bias"].numpy()).reshape(
                  bias_shape
              )
          )
      else:
          objax_layer.b = None
  
imports: import functools
  from typing import Callable, Sequence, Union, Optional
  from jax import numpy
  import objax
  from objax.constants import ConvPadding
  from objax.nn import Conv2D
  from objax.typing import JaxArray, ConvPaddingInt
label_desc: Use image node_id 3 for calling a global function with example usage: convert_conv_or_linear_layer(keras_layer, objax_layer) without return types