code: def extract(a, t, x_shape):
      b, *_ = t.shape
      out = a.gather(-1, t)
      return out.reshape(b, *((1,) * (len(x_shape) - 1)))
  
imports: from functools import partial
  from inspect import isfunction
  import numpy
  import torch
  from torch import nn, einsum
  import torch.nn.functional
label_desc: Use image node_id 2 for calling a global function with example usage: extract(a, t, x_shape) and returns: out