def METHOD_NAME(base_layer_fn, input_shape, data_init):
    sample_data = np.ones([1] + input_shape, dtype=np.float32)
    base_layer = base_layer_fn()
    wn_layer = wrappers.WeightNormalization(base_layer, data_init)
    wn_output = wn_layer(sample_data)
    wn_removed_layer = wn_layer.remove()
    wn_removed_output = wn_removed_layer(sample_data)
    np.testing.assert_allclose(wn_removed_output.numpy(), wn_output.numpy())
    assert isinstance(wn_removed_layer, base_layer.__class__)