Commit 9354ab24 authored by Mingsheng Hong's avatar Mingsheng Hong Committed by TensorFlower Gardener
Browse files

Mark the If op stateful, so that it won't get pruned away in a graph function,

when we only run that If op for its side effects (e.g. enqueuing tensors.)

Also extended the kernel impl to handle the case where the kernel works with
multiple functional libs through its lifetime (b/37549631). The code is modeled
after the WhileOp kernel impl.

An example graph function that runs If for its side-effects is:

function {
    signature {
      name: "S12control_flow23testTensorEnqueueInCondyySb_SftF.tf_CPU.device_partition"
      input_arg {
        name: "arg_0"
        type: DT_FLOAT # DT_BOOL
      }
    }
    node_def {
      name: "op/testTensorEnqueueInCond.14.14"
      op: "Const"
      device: "/device:CPU:0"
      attr {
        key: "dtype"
        value {
          type: DT_FLOAT
        }
      }
      attr {
        key: "value"
        value {
          tensor {
            dtype: DT_FLOAT
            tensor_shape {
            }
            float_val: 1
          }
        }
      }
    }
    node_def {
      name: "op/testTensorEnqueueInCond_5.22.3"
      op: "If"
      input: "arg_0"
      input: "op/testTensorEnqueueInCond.14.14:output:0"
      attr {
        key: "Tcond"
        value {
          type: DT_FLOAT # DT_BOOL
        }
      }
      attr {
        key: "Tin"
        value {
          list {
            type: DT_FLOAT
          }
        }
      }
      attr {
        key: "Tout"
        value {
          list {
          }
        }
      }
      attr {
        key: "else_branch"
        value {
          func {
            name: "false/testTensorEnqueueInCond_4.22.3"
          }
        }
      }
      attr {
        key: "then_branch"
        value {
          func {
            name: "true/testTensorEnqueueInCond_3.22.3"
          }
        }
      }
    }
  }

PiperOrigin-RevId: 206983563
parent 68a5b035
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment