reshape = tf.reshape(max_pool2, [-1, flatten_size])

    # output shape: [-1, 1024]
    fc1 = gen_mnist_ops.new_fully_connected(reshape, self.weights['f3'],
                                            self.biases['b3'], 'RELU')
    # output shape: [-1, 10]
    return gen_mnist_ops.new_fully_connected(fc1, self.weights['f4'],
                                             self.biases['b4'])


def main(strategy):

    bias = tf.zeros([3])
    kwargs = {'input_': input_, 'filter_': filter_, 'bias': bias, 'act': 'RELU'}

    self._assertOpAndComposite([input_, filter_, bias],
                               tf.function(gen_mnist_ops.new_fully_connected),
                               ops_defs._composite_fully_connected, kwargs)

  def test_new_max_pool(self):
    input_ = tf.random.uniform([8, 4, 4, 1])
    kwargs = {
        'input_': input_,

          shape_1,
          grad,
          strides=strides,
          padding=padding,
          dilations=dilations,
          data_format='NHWC'), bias_grad
  ]


@Composite(
    'NewFullyConnected',
    inputs=['input_: T', 'filter_: T', 'bias: T'],
    attrs=['act: {"", "RELU", "RELU6", "TANH"} = ""'],
    derived_attrs=['T: {float, int8}'],
    outputs=['o: T'])