code: def differentiable_barycentrics(
      framebuffer: fb.Framebuffer,
      clip_space_vertices: type_alias.TensorLike,
      triangles: type_alias.TensorLike,
      use_vectorized_map: bool = True,
  ) -> fb.Framebuffer:
      """Computes differentiable barycentric coordinates from a Framebuffer.
  
      The barycentric coordinates will be differentiable w.r.t. the input vertices.
      Later, we may support derivatives w.r.t. pixel position for mip-mapping.
  
      Args:
        framebuffer: a multi-layer Framebuffer containing triangle ids and a
          foreground mask with shape [batch, num_layers, height, width, 1]
        clip_space_vertices: a 2-D float32 tensor with shape [vertex_count, 4] or a
          3-D tensor with shape [batch, vertex_count, 4] containing homogenous
          vertex positions (xyzw).
        triangles: a 2-D int32 tensor with shape [triangle_count, 3] or a 3-D tensor
          with shape [batch, triangle_count, 3] containing per-triangle vertex
          indices in counter-clockwise order.
        use_vectorized_map: If true uses vectorized_map otherwise uses map_fn.
  
      Returns:
        a copy of `framebuffer`, but the differentiable barycentric coordinates will
        replace any barycentric coordinates already in the `framebuffer`.
      """
      rank = lambda t: len(t.shape)
  
      clip_space_vertices = tf.convert_to_tensor(clip_space_vertices)
      shape.check_static(
          tensor=clip_space_vertices,
          tensor_name="clip_space_vertices",
          has_rank_greater_than=1,
          has_rank_less_than=4,
      )
      if rank(clip_space_vertices) == 2:
          clip_space_vertices = tf.expand_dims(
              clip_space_vertices, axis=0
          )
  
      triangles = tf.convert_to_tensor(triangles)
      shape.check_static(
          tensor=triangles,
          tensor_name="triangles",
          has_rank_greater_than=1,
          has_rank_less_than=4,
      )
      if rank(triangles) == 2:
          triangles = tf.expand_dims(triangles, axis=0)
      else:
          shape.compare_batch_dimensions(
              tensors=(clip_space_vertices, triangles),
              last_axes=(-3, -3),
              broadcast_compatible=False,
          )
  
      shape.compare_batch_dimensions(
          tensors=(clip_space_vertices, framebuffer.triangle_id),
          last_axes=(-3, -4),
          broadcast_compatible=False,
      )
  
      # Compute image pixel coordinates.
      px, py = normalized_pixel_coordinates(
          framebuffer.width, framebuffer.height
      )
  
      def compute_barycentrics_fn(
          slices: Tuple[
              type_alias.TensorLike,
              type_alias.TensorLike,
              type_alias.TensorLike,
          ]
      ) -> tf.Tensor:
          (
              clip_vertices_slice,
              triangle_slice,
              triangle_id_slice,
          ) = slices
          triangle_id_slice = triangle_id_slice[..., 0]
          if (
              rank(triangle_id_slice) == 2
          ):  # There is no layer dimension.
              triangle_id_slice = tf.expand_dims(
                  triangle_id_slice, axis=0
              )
          # Compute per-triangle inverse matrices.
          triangle_matrices = compute_triangle_matrices(
              clip_vertices_slice, triangle_slice
          )
  
          # Compute per-pixel barycentric coordinates.
          barycentric_coords = compute_barycentric_coordinates(
              triangle_id_slice, triangle_matrices, px, py
          )
          barycentric_coords = tf.transpose(
              barycentric_coords, perm=[1, 2, 3, 0]
          )
          return barycentric_coords
  
      if use_vectorized_map:
          per_image_barycentrics = tf.vectorized_map(
              compute_barycentrics_fn,
              (clip_space_vertices, triangles, framebuffer.triangle_id),
          )
      else:
          num_meshes = tf.shape(clip_space_vertices)[0]
          triangles_repeated = tf.repeat(
              triangles, repeats=num_meshes, axis=0
          )
          per_image_barycentrics = tf.map_fn(
              compute_barycentrics_fn,
              (
                  clip_space_vertices,
                  triangles_repeated,
                  framebuffer.triangle_id,
              ),
              fn_output_signature=tf.TensorSpec(
                  shape=(1, None, None, 3)
              ),
          )
  
      barycentric_coords = tf.stack(per_image_barycentrics, axis=0)
      # After stacking barycentrics will have layers dimension no matter what.
      # In order to make sure we return differentiable barycentrics of the same
      # shape - reshape the tensor using original shape.
      barycentric_coords = tf.reshape(
          barycentric_coords,
          shape=tf.shape(framebuffer.barycentrics.value),
      )
      # Mask out barycentrics for background pixels.
      barycentric_coords = (
          barycentric_coords * framebuffer.foreground_mask
      )
  
      return fb.Framebuffer(
          triangle_id=framebuffer.triangle_id,
          vertex_ids=framebuffer.vertex_ids,
          foreground_mask=framebuffer.foreground_mask,
          attributes=framebuffer.attributes,
          barycentrics=fb.RasterizedAttribute(
              barycentric_coords, None, None
          ),
      )
  
file_path: graphics/tensorflow_graphics/rendering/barycentrics.py
imports: from typing import Tuple
  import tensorflow
  from tensorflow_graphics.rendering import framebuffer
  from tensorflow_graphics.util import shape
  from tensorflow_graphics.util import type_alias
label_desc: Use image node_id 1 for calling a global function with example usage: differentiable_barycentrics(framebuffer, clip_space_vertices, triangles, use_vectorized_map) and returns: fb, barycentric_coords