code: def _slice(x, begin, size, axis):
      """Slice wrapper which has XLA-compatible gradients."""
  
      # Not in an XLA context => don't need a custom gradient, can use pad.
      if (
          tf.executing_eagerly()
          or not control_flow_util.GraphOrParentsInXlaContext(
              tf1.get_default_graph()
          )
      ):
          return tf.slice(x, begin, size)
  
      # TODO(b/143313126): Upstream something like this to TF proper.
      @tf.custom_gradient
      def _custom_slice_helper(x):
          """tf.slice uses tf.pad for backprop, unknown shape breaks XLA."""
  
          def grad(dy):
              dshp = tf.shape(x) - tf.shape(dy)
              z = tf.zeros(
                  tf.where(tf.equal(0, dshp), tf.shape(x), dshp),
                  dtype=x.dtype,
              )
              dx = tf.roll(
                  tf.concat([dy, z], axis=axis),
                  shift=begin[axis],
                  axis=axis,
              )
              return dx
  
          return tf.slice(x, begin, size), grad
  
      return _custom_slice_helper(x)
  
imports: import functools
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 1 for calling a global function with example usage: _slice(x, begin, size, axis) and returns: _custom_slice_helper, tf, tf,  grad, dx