# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # where

import torch

judge = torch.rand(2,2)
judge

a = torch.zeros(2,2)
b = torch.ones(2,2)

torch.where(judge>0.5, a, b)

# # gather

norm = torch.randn(3,4)
norm

ind = norm.topk(k = 2)[1]
ind

ex_gather = torch.arange(4)+10
ex_gather.expand(3,4)

torch.gather(ex_gather.expand(3,4), dim = 1, index = ind)

StarcoderdataJupyterScriptsDedupFiltered