tf.switch_case

Создайте операцию switch/case, т.е. условную операцию с целочисленным индексом.

tf.switch_case(
    branch_index, branch_fns, default=None, name='switch_case'
)

См. также tf.case.

Эта операция может быть значительно эффективнее, чем tf.case, когда будет выбран ровно один ответвление. tf.switch_case больше похожа на оператор switch/case в C++, чем tf.case, которая больше похожа на цепочку if/elif/elif/else.

Параметр branch_fns — это либо список пар (int, вызываемый объект), либо просто список вызываемых объектов (в этом случае индекс неявно является ключом). branch_index Tensor используется для выбора элемента в branch_fns с совпадающим ключом int, возвращаясь к default, если ничего не совпадает, или max(keys), если default не указан. Ключи должны образовывать непрерывный набор от 0 до len(branch_fns) - 1.

tf.switch_case поддерживает вложенные структуры, как реализовано в tf.nest. Все вызываемые объекты должны возвращать одну и ту же (возможно, вложенную) структуру списков, кортежей и/или именованных кортежей.

Пример:

Псевдокод:

switch (branch_index) {  // c-style switch
  case 0: return 17;
  case 1: return 31;
  default: return -1;
}

или

branches = {0: lambda: 17, 1: lambda: 31}
branches.get(branch_index, lambda: -1)()

Выражения:

def f1(): return tf.constant(17)
def f2(): return tf.constant(31)
def f3(): return tf.constant(-1)
r = tf.switch_case(branch_index, branch_fns={0: f1, 1: f2}, default=f3)
# Equivalent: tf.switch_case(branch_index, branch_fns={0: f1, 1: f2, 2: f3})

Совместимость с V2

branch_fns в версии 1 мог быть словарем. Однако в версии 2 tf.Tensor и tf.Variable больше нельзя использовать в качестве ключей словаря, поэтому их нельзя использовать в качестве ключей для словаря. Используйте вместо них список или кортеж.