tf2torch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TF to Torch Conversion\n",
    "\n",
    "## Purpose\n",
    "- Weight transfer\n",
    "- Torch reproducibility\n",
    "- Torch debugging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "p6kExhij7_gs",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import torch\n",
    "import tensorflow as tf\n",
    "import gdown\n",
    "import tqdm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import neurite as ne\n",
    "import voxelmorph as vxm\n",
    "\n",
    "from torch import nn\n",
    "from tensorflow.keras.models import Model\n",
    "from skimage.metrics import structural_similarity\n",
    "\n",
    "# local code\n",
    "from synthmorph import layers, networks, datamodule as dm, utils\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu' # note: only gpu has been tested so far\n",
    "torch.multiprocessing.set_start_method('spawn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Prevent TF model from taking whole GPU memory\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "for gpu in gpus:\n",
    "    tf.config.experimental.set_memory_growth(gpu, True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deformable Registration\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transfer Weights TF -> Torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define fresh Keras model, only for loading original author's weights\n",
    "# This section is just a copy of the orginal demo to define the Keras model\n",
    "\n",
    "# Label maps\n",
    "in_shape = (256,) * 2\n",
    "num_dim = len(in_shape)\n",
    "num_label = 16\n",
    "num_maps = 1\n",
    "label_maps = []\n",
    "for _ in range(num_maps):\n",
    "    # Draw image and warp.\n",
    "    im = ne.utils.augment.draw_perlin(\n",
    "        out_shape=(*in_shape, num_label),\n",
    "        scales=(32, 64), max_std=1,\n",
    "    )\n",
    "    warp = ne.utils.augment.draw_perlin(\n",
    "        out_shape=(*in_shape, num_label, num_dim),\n",
    "        scales=(16, 32, 64), max_std=16,\n",
    "    )\n",
    "\n",
    "    # Transform and create label map.\n",
    "    im = vxm.utils.transform(im, warp)\n",
    "    lab = tf.argmax(im, axis=-1)\n",
    "    label_maps.append(np.uint8(lab))\n",
    "\n",
    "# Image generator\n",
    "gen_arg = dict(\n",
    "    in_shape=in_shape,\n",
    "    in_label_list=np.unique(label_maps),\n",
    "    warp_std=3,\n",
    "    warp_res=(8, 16, 32),\n",
    ")\n",
    "gen_model_1 = ne.models.labels_to_image(**gen_arg, id=1)\n",
    "gen_model_2 = ne.models.labels_to_image(**gen_arg, id=2)\n",
    "\n",
    "# Registration model.\n",
    "reg_model = vxm.networks.VxmDense(\n",
    "    inshape=in_shape,\n",
    "    int_resolution=2,\n",
    "    svf_resolution=2,\n",
    "    nb_unet_features=([256] * 4, [256] * 8),\n",
    "    reg_field='warp',\n",
    ")\n",
    "\n",
    "# Model for optimization.\n",
    "ima_1, map_1 = gen_model_1.outputs\n",
    "ima_2, map_2 = gen_model_2.outputs\n",
    "\n",
    "_, warp = reg_model((ima_1, ima_2))\n",
    "pred = vxm.layers.SpatialTransformer(fill_value=0)((map_1, warp))\n",
    "\n",
    "inputs = gen_model_1.inputs + gen_model_2.inputs\n",
    "out = (map_2, pred)\n",
    "model = tf.keras.Model(inputs, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load Keras pretrained weights\n",
    "gdown.download('https://drive.google.com/uc?id=1xridvtyEWgWsWJPYVrQfDCtSgbj2beRz', 'weights.h5')\n",
    "model.load_weights('weights.h5')\n",
    "\n",
    "# Extract weights from the registration model only\n",
    "keras_vxmdense = reg_model   \n",
    "keras_weights = {w.name: (w.numpy(), w.dtype, w.shape) for w in keras_vxmdense.weights}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Only get kernel weight (ie. skip biases)\n",
    "keras_weights_keys = list(keras_weights.keys())\n",
    "keras_kernels = [string for string in keras_weights_keys if 'bias' not in string]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define fresh Torch model\n",
    "vol_size = (256,) * 2\n",
    "unet_enc_nf = [256] * 4\n",
    "unet_dec_nf = [256] * 8\n",
    "int_steps = 7 \n",
    "int_downsize = 2\n",
    "bidir = False\n",
    "torch_vxmdense = networks.VxmDense(\n",
    "    inshape=vol_size,\n",
    "    nb_unet_features=[unet_enc_nf, unet_dec_nf],\n",
    "    int_steps=int_steps,\n",
    "    int_downsize=int_downsize,\n",
    "    bidir=bidir,\n",
    "    unet_half_res=True,\n",
    ")\n",
    "\n",
    "torch_weights = torch_vxmdense.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_weights = {}\n",
    "\n",
    "# Transfer the weights (the order of layers are the same)\n",
    "for k,t in zip(keras_weights.keys(), torch_weights.keys()):\n",
    "    if k in keras_kernels:\n",
    "        new_weights[t] = torch.Tensor(np.moveaxis(keras_weights[k][0], [-1, -2], [0, 1]))\n",
    "    else:\n",
    "        new_weights[t] = torch.Tensor(keras_weights[k][0])\n",
    "\n",
    "torch_vxmdense.load_state_dict(new_weights)\n",
    "# torch.save(torch_vxmdense.state_dict(), Path(\".\") / 'authors.pth' )  # uncomment to save weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PyTorch Reimplementation Debug\n",
    "Make sure that both TF and Torch models are using the same weights "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Data preprocessing for TF\n",
    "def tf_conform(x, in_shape=in_shape):\n",
    "    '''Resize and normalize image.'''\n",
    "    x = np.float32(x)\n",
    "    x = np.squeeze(x)\n",
    "    x = ne.utils.minmax_norm(x)\n",
    "    x = ne.utils.zoom(x, zoom_factor=[o / i for o, i in zip(in_shape, x.shape)])\n",
    "    return np.expand_dims(x, axis=(0, -1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load MNIST in TF\n",
    "images, digits = tf.keras.datasets.mnist.load_data()[-1]\n",
    "ind = np.flatnonzero(digits == 6)\n",
    "moving = tf_conform(images[ind[256]])\n",
    "fixed = tf_conform(images[ind[22]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Data preprocessing for Torch\n",
    "torch_conform = lambda x, size: dm.conform(x, size, device) \n",
    "torch_moving = torch_conform(moving, (256,256))\n",
    "torch_fixed = torch_conform(fixed, (256,256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Custom models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load TF layers by indexing\n",
    "custom_keras_layers = keras_vxmdense.layers[:]\n",
    "custom_keras_model= Model(\n",
    "    inputs=keras_vxmdense.inputs, \n",
    "    outputs=custom_keras_layers[-1].output    # output from chosen layer\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load Torch layers by specifying each module\n",
    "# Note: SpatialTransformer is not compatible with nn.Sequential, hence separated\n",
    "unet_torch = torch_vxmdense.unet_model\n",
    "flow_torch = torch_vxmdense.flow\n",
    "vecint_torch = torch_vxmdense.integrate\n",
    "rescale_torch = torch_vxmdense.fullsize\n",
    "spatial_torch = torch_vxmdense.transformer\n",
    "custom_torch_model = nn.Sequential(\n",
    "    unet_torch, \n",
    "    flow_torch,\n",
    "    vecint_torch,\n",
    "    rescale_torch,\n",
    ")\n",
    "custom_torch_model = custom_torch_model.eval().cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### SSIM of TF vs Torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare chosen output from TF and Torch registration models\n",
    "keras_output = custom_keras_model.predict((moving, fixed))\n",
    "keras_output = keras_output.transpose(0, 3, 1, 2).squeeze(0)\n",
    "\n",
    "torch_input = torch.cat([torch_moving, torch_fixed], dim=1)\n",
    "torch_output = custom_torch_model(torch_input)\n",
    "torch_output = spatial_torch([torch_moving, torch_output]) # uncomment only when comparing moved image (i.e. whole reg model)\n",
    "torch_output = torch_output.squeeze(0).cpu().detach().numpy()\n",
    "\n",
    "data_range = torch_output.max() - torch_output.min()\n",
    "channel_axis = 0 if torch_output.ndim > 2 else None\n",
    "ssim_mean, ssim_full= structural_similarity(\n",
    "    torch_output, \n",
    "    keras_output,\n",
    "    win_size=11,    # must be odd\n",
    "    data_range=data_range,\n",
    "    channel_axis=0,\n",
    "    multichannel=True,\n",
    "    full=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot each axis of resulting SSIM\n",
    "num_plots = ssim_full.shape[0]\n",
    "plot_scale = 6\n",
    "fig, axs = plt.subplots(1, num_plots, figsize=(plot_scale*num_plots, plot_scale), squeeze=False)   # subplots in  one row\n",
    "fig.suptitle(f\"SSIM Plot for each channel, Mean = {ssim_mean*100:.4f}\")\n",
    "for i in range(num_plots):\n",
    "    ax = axs[0, i]\n",
    "    im = ax.imshow(ssim_full[i], cmap='gray')\n",
    "    ax.set_title(f'Channel {i+1}')\n",
    "    ax.get_xaxis().set_visible(False)\n",
    "    ax.get_yaxis().set_visible(False)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inference time TF vs PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(seed=42)\n",
    "width, height = 256, 256\n",
    "channel = 1\n",
    "batch_size = 1\n",
    "shape = (batch_size, width, height, channel)\n",
    "moving = rng.standard_normal(size=shape)\n",
    "fixed = rng.standard_normal(size=shape)\n",
    "\n",
    "prepare_torch = lambda x: torch.from_numpy(x).to(device, torch.float32).permute(0, -1, 1, 2)\n",
    "torch_moving = prepare_torch(moving)\n",
    "torch_fixed = prepare_torch(fixed)\n",
    "\n",
    "tf_moving = tf.constant(moving)\n",
    "tf_fixed = tf.constant(fixed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 100 -r 10 -p 4 \n",
    "keras_vxmdense.predict((tf_moving, tf_fixed), verbose=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 100 -r 10 -p 4\n",
    "torch_vxmdense(torch_moving, torch_fixed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Affine Registration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transfer weights TF -> Torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Keras pretrained weights\n",
    "keras_weights_path = Path(\".\") / \"weights\" / \"keras\"\n",
    "keras_affine_path = keras_weights_path / \"affine_author_2d_256.h5\"\n",
    "# gdown.download('https://drive.google.com/uc?id=1DWiVxCvQmYDSBS1RcVbUeTxX0XArQsGv', affine_weights)\n",
    "\n",
    "in_shape = (256,) * 2\n",
    "keras_model= vxm.networks.VxmAffineFeatureDetector(in_shape)\n",
    "keras_model.load_weights(keras_affine_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only get kernel weights (ie. skip biases)\n",
    "keras_weights = {w.name: (w.numpy(), w.dtype, w.shape) for w in keras_model.weights}\n",
    "keras_weights_keys = list(keras_weights.keys())\n",
    "keras_kernels = [string for string in keras_weights_keys if 'bias' not in string]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Torch model and weigths\n",
    "torch_model = networks.VxmAffineFeatureDetector(in_shape)\n",
    "torch_weights = torch_model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_weights = {}\n",
    "# Transfer the weights (the order of layers are the same)\n",
    "for k,t in zip(keras_weights.keys(), torch_weights.keys()):\n",
    "    if k in keras_kernels:\n",
    "        new_weights[t] = torch.Tensor(np.moveaxis(keras_weights[k][0], [-1, -2], [0, 1]))\n",
    "    else:\n",
    "        new_weights[t] = torch.Tensor(keras_weights[k][0])\n",
    "\n",
    "torch_model.load_state_dict(new_weights)\n",
    "\n",
    "torch_weights_path = Path(\".\") / \"weights\" / \"torch\"\n",
    "torch_affine_path = torch_weights_path / \"affine_author_2d_256.pth\"\n",
    "torch.save(torch_model.state_dict(), torch_affine_path)  # uncomment to save weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Debug Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng()\n",
    "ssim = lambda x, y, drange: structural_similarity(\n",
    "    x, y,\n",
    "    win_size=11,    # must be odd\n",
    "    data_range=drange,\n",
    "    channel_axis=0,\n",
    "    multichannel=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_path = Path(\".\") / \"weights\"\n",
    "keras_weights_path = weights_path / \"keras\"\n",
    "keras_affine_path = keras_weights_path / \"affine_author_2d_256.h5\"\n",
    "torch_weights_path = weights_path / \"torch\"\n",
    "torch_affine_path = torch_weights_path / \"affine_author_2d_256.pth\"\n",
    "\n",
    "in_shape = (256,) * 2\n",
    "keras_model = vxm.networks.VxmAffineFeatureDetector(in_shape)\n",
    "keras_model.load_weights(keras_affine_path)\n",
    "torch_model = networks.VxmAffineFeatureDetector(in_shape)\n",
    "torch_model.load_state_dict(torch.load(torch_affine_path))\n",
    "torch_model = torch_model.eval()\n",
    "\n",
    "# Load TF layers by indexing\n",
    "custom_keras_layers = keras_model.layers[:]\n",
    "custom_keras_model= Model(\n",
    "    inputs=keras_vxmdense.inputs, \n",
    "    outputs=custom_keras_layers[-1].output    # output from chosen layer\n",
    ")\n",
    "\n",
    "torch_unet = torch_model.det\n",
    "custom_torch_model = nn.Sequential(\n",
    "    torch_unet,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feat_1 = rng.random((1, 8, 8, 64), dtype=np.float32)\n",
    "feat_2 = rng.random((1, 8, 8, 64), dtype=np.float32)\n",
    "tf_feat_1 = tf.constant(feat_1)\n",
    "torch_feat_1 = torch.tensor(feat_1)\n",
    "tf_feat_2 = tf.constant(feat_2)\n",
    "torch_feat_2 = torch.tensor(feat_2)\n",
    "\n",
    "# Center of Mass (barycenter)\n",
    "shape_full = (256, 256)\n",
    "num_dim = len(shape_full)\n",
    "args = dict(axes=range(1, num_dim + 1), normalize=True, shift_center=True)\n",
    "tf_com = lambda x: ne.utils.barycenter(x, **args)\n",
    "torch_com = lambda x: utils.barycenter(x, **args)\n",
    "\n",
    "torch_cen_1 = torch_com(torch_feat_1) * torch.tensor(shape_full)\n",
    "tf_cen_1 = tf_com(tf_feat_1) * shape_full\n",
    "torch_cen_2 = torch_com(torch_feat_2) * torch.tensor(shape_full)\n",
    "tf_cen_2 = tf_com(tf_feat_2) * shape_full\n",
    "# print(torch_cen_1 - tf_cen_2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run previous cell first\n",
    "# Least squares fit\n",
    "\n",
    "# Channel weights for WLS\n",
    "axes = tuple(range(1, num_dim + 1))\n",
    "pow_1 = np.sum(feat_1, axis=axes)\n",
    "pow_2 = np.sum(feat_2, axis=axes)\n",
    "pow_1 /= np.sum(pow_1, axis=-1, keepdims=True)\n",
    "pow_2 /= np.sum(pow_2, axis=-1, keepdims=True)\n",
    "lsweights = pow_1 * pow_2\n",
    "\n",
    "tf_lsweights = tf.constant(lsweights)\n",
    "torch_lsweights = torch.tensor(lsweights)\n",
    "\n",
    "tf_lsfit = lambda x, y, weights=None: vxm.utils.fit_affine(x, y, weights)\n",
    "torch_lsfit = lambda x, y, weights=None: utils.fit_affine(x, y, weights)\n",
    "\n",
    "tf_aff = tf_lsfit(tf_cen_1, tf_cen_2, tf_lsweights)\n",
    "torch_aff = torch_lsfit(torch_cen_1, torch_cen_2, torch_lsweights)\n",
    "\n",
    "print(f\"torch_aff: {torch_aff}\\ntf_aff: {tf_aff}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compose transforms\n",
    "tf_compose = lambda x: vxm.layers.ComposeTransform(shift_center=False)(x)\n",
    "torch_compose = lambda x: layers.ComposeTransform(shift_center=False)(x)\n",
    "shape_full = np.asanyarray((256, 256))\n",
    "num_dim = 2\n",
    "batch_size= feat_1.shape[0]\n",
    "def tensor(x):\n",
    "    x = np.asarray(x[None, :-1, :], )\n",
    "    return np.repeat(x, repeats=batch_size, axis=0).astype(np.float32)\n",
    "\n",
    "def cen(shape):\n",
    "    mat = np.eye(num_dim + 1)\n",
    "    mat[:-1, -1] = -0.5 * (shape - 1)\n",
    "    return tensor(mat)\n",
    "\n",
    "def un_cen(shape):\n",
    "    mat = np.eye(num_dim + 1)\n",
    "    mat[:-1, -1] = +0.5 * (shape - 1)\n",
    "    return tensor(mat)\n",
    "\n",
    "def scale(fact):\n",
    "    mat = np.diag((*[fact] * num_dim, 1))\n",
    "    return tensor(mat)\n",
    "\n",
    "np_cen = cen(shape_full)\n",
    "np_uncen = un_cen(shape_full)\n",
    "tf_cen, tf_uncen = tf.constant(np_cen), tf.constant(np_uncen)\n",
    "torch_cen, torch_uncen = torch.tensor(np_cen), torch.tensor(np_uncen)\n",
    "print(tf_uncen.shape)\n",
    "tf_composed = tf_compose([tf_uncen, tf_aff, tf_cen])\n",
    "torch_composed = torch_compose([torch_uncen, torch_aff, torch_cen])\n",
    "\n",
    "print(f\"tf_composed: {tf_composed}\\ntorch_composed: {torch_composed}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run previous cells \n",
    "# Affine to dense shift\n",
    "shape_out = (256, 256)\n",
    "tf_aff2dense = lambda x: vxm.utils.affine_to_dense_shift(x, shape_out, False)\n",
    "torch_aff2dense = lambda x: utils.affine_to_dense_shift(x, shape_out, False)\n",
    "\n",
    "tf_dense = tf_aff2dense(tf_aff)\n",
    "torch_dense = torch_aff2dense(torch_aff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TF to Torch functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### tf.gather(params, indices, axis, batch_dims)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}