Switches between two operations depending on a scalar value. Note that both `then_expression` and `else_expression` should be symbolic tensors of the *same shape*. Arguments: condition: tensor (`int` or `bool`). then_expression: either a tensor, or a callable that returns a tenso
(condition, then_expression, else_expression)
| 4079 | |
| 4080 | @keras_export('keras.backend.switch') |
| 4081 | def switch(condition, then_expression, else_expression): |
| 4082 | """Switches between two operations depending on a scalar value. |
| 4083 | |
| 4084 | Note that both `then_expression` and `else_expression` |
| 4085 | should be symbolic tensors of the *same shape*. |
| 4086 | |
| 4087 | Arguments: |
| 4088 | condition: tensor (`int` or `bool`). |
| 4089 | then_expression: either a tensor, or a callable that returns a tensor. |
| 4090 | else_expression: either a tensor, or a callable that returns a tensor. |
| 4091 | |
| 4092 | Returns: |
| 4093 | The selected tensor. |
| 4094 | |
| 4095 | Raises: |
| 4096 | ValueError: If rank of `condition` is greater than rank of expressions. |
| 4097 | """ |
| 4098 | if condition.dtype != dtypes_module.bool: |
| 4099 | condition = math_ops.cast(condition, 'bool') |
| 4100 | cond_ndim = ndim(condition) |
| 4101 | if not cond_ndim: |
| 4102 | if not callable(then_expression): |
| 4103 | |
| 4104 | def then_expression_fn(): |
| 4105 | return then_expression |
| 4106 | else: |
| 4107 | then_expression_fn = then_expression |
| 4108 | if not callable(else_expression): |
| 4109 | |
| 4110 | def else_expression_fn(): |
| 4111 | return else_expression |
| 4112 | else: |
| 4113 | else_expression_fn = else_expression |
| 4114 | x = control_flow_ops.cond(condition, then_expression_fn, else_expression_fn) |
| 4115 | else: |
| 4116 | # tf.where needs its condition tensor |
| 4117 | # to be the same shape as its two |
| 4118 | # result tensors |
| 4119 | if callable(then_expression): |
| 4120 | then_expression = then_expression() |
| 4121 | if callable(else_expression): |
| 4122 | else_expression = else_expression() |
| 4123 | expr_ndim = ndim(then_expression) |
| 4124 | if cond_ndim > expr_ndim: |
| 4125 | raise ValueError('Rank of `condition` should be less than or' |
| 4126 | ' equal to rank of `then_expression` and ' |
| 4127 | '`else_expression`. ndim(condition)=' + str(cond_ndim) + |
| 4128 | ', ndim(then_expression)' |
| 4129 | '=' + str(expr_ndim)) |
| 4130 | if cond_ndim > 1: |
| 4131 | ndim_diff = expr_ndim - cond_ndim |
| 4132 | cond_shape = array_ops.concat( |
| 4133 | [array_ops.shape(condition), [1] * ndim_diff], axis=0) |
| 4134 | condition = array_ops.reshape(condition, cond_shape) |
| 4135 | expr_shape = array_ops.shape(then_expression) |
| 4136 | shape_diff = expr_shape - cond_shape |
| 4137 | tile_shape = array_ops.where(shape_diff > 0, expr_shape, |
| 4138 | array_ops.ones_like(expr_shape)) |