Python Flatten Examples

Example #1

File: combinator_test.py Project: kchoworks/tensorflow-probability

def define_cnn():
  return combinator.Serial([convolution.Conv(20, (2, 2)),
                            normalization.BatchNorm(),
                            core.Relu(),
                            reshape.Flatten(),
                            core.Dense(10),
                            core.Softmax()])

Example #2

File: reshape_test.py Project: kchoworks/tensorflow-probability

    def test_flatten_call(self):
        net_rng, data_rng = random.split(random.PRNGKey(0))

        net_init = reshape.Flatten()

        net = net_init.init(net_rng, state.Shape((10, 10)))
        x = random.normal(data_rng, (1, 10, 10))
        self.assertEqual(jax.vmap(net)(x).shape, (1, 100))

        x = random.normal(data_rng, (5, 20, 5))
        self.assertEqual(jax.vmap(net)(x).shape, (5, 100))

Example #3

File: reshape_test.py Project: kchoworks/tensorflow-probability

    def test_flatten_shape(self):
        net_init = reshape.Flatten()

        out_shape = net_init.spec(state.Shape((5, 100))).shape
        self.assertEqual(out_shape, (500, ))

        out_shape = net_init.spec(state.Shape((10, 10))).shape
        self.assertEqual(out_shape, (100, ))

        out_shape = net_init.spec(state.Shape((1, 2, 5, 10))).shape
        self.assertEqual(out_shape, (100, ))