code: def WeightedSum():
      """Returns a layer that computes a weighted sum of the given values."""
  
      def f(values, weights):  # pylint: disable=invalid-name
          return jnp.sum(values * weights)
  
      return base.Fn("WeightedSum", f)
  
imports: from trax import fastmath
  from trax import shapes
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
label_desc: Use image node_id 20 for calling a global function with example usage: WeightedSum() and returns: base, jnp