<s> import unittest <EOL> import numpy <EOL> import chainer <EOL> from chainer import cuda <EOL> from chainer import functions <EOL> from chainer import gradient_check <EOL> from chainer import testing <EOL> from chainer . testing import attr <EOL> @ testing . parameterize ( <EOL> { '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:7> , <NUM_LIT:3> ) , '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : [ <NUM_LIT:2> , <NUM_LIT:5> ] , <EOL> '<STR_LIT>' : [ [ slice ( None ) , slice ( None , <NUM_LIT:2> ) ] , [ slice ( None ) , slice ( <NUM_LIT:2> , <NUM_LIT:5> ) ] , <EOL> [ slice ( None ) , slice ( <NUM_LIT:5> , None ) ] ] } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:7> , <NUM_LIT:3> ) , '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : [ <NUM_LIT:2> , <NUM_LIT:5> ] , <EOL> '<STR_LIT>' : [ slice ( None , <NUM_LIT:2> ) , slice ( <NUM_LIT:2> , <NUM_LIT:5> ) , slice ( <NUM_LIT:5> , None ) ] } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:9> , <NUM_LIT:3> ) , '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:3> , <EOL> '<STR_LIT>' : [ [ slice ( None ) , slice ( None , <NUM_LIT:3> ) ] , [ slice ( None ) , slice ( <NUM_LIT:3> , <NUM_LIT:6> ) ] , <EOL> [ slice ( None ) , slice ( <NUM_LIT:6> , None ) ] ] } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:2> , <NUM_LIT:6> , <NUM_LIT:3> ) , '<STR_LIT>' : <NUM_LIT:1> , '<STR_LIT>' : <NUM_LIT:3> , <EOL> '<STR_LIT>' : [ [ slice ( None ) , slice ( None , <NUM_LIT:2> ) ] , [ slice ( None ) , slice ( <NUM_LIT:2> , <NUM_LIT:4> ) ] , <EOL> [ slice ( None ) , slice ( <NUM_LIT:4> , None ) ] ] } , <EOL> { '<STR_LIT>' : ( <NUM_LIT:2> , ) , '<STR_LIT>' : <NUM_LIT:0> , '<STR_LIT>' : [ <NUM_LIT:1> ] , <EOL> '<STR_LIT>' : [ slice ( None , <NUM_LIT:1> ) , slice ( <NUM_LIT:1> , None ) ] } , <EOL> ) <EOL> class TestSplitAxis ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . x = numpy . arange ( <EOL> numpy . prod ( self . shape ) , dtype = numpy . float32 ) . reshape ( self . shape ) <EOL> self . ys = [ self . x [ s ] for s in self . slices ] <EOL> def check_forward ( self , x_data , ys_data , indices_or_sections , axis ) : <EOL> x = chainer . Variable ( x_data ) <EOL> ys = functions . split_axis ( x , indices_or_sections , axis ) <EOL> for yd , y in zip ( ys_data , ys ) : <EOL> self . assertIsInstance ( y . data . shape , tuple ) <EOL> gradient_check . assert_allclose ( yd , y . data , atol = <NUM_LIT:0> , rtol = <NUM_LIT:0> ) <EOL> def test_forward_cpu ( self ) : <EOL> self . check_forward ( self . x , self . ys , self . ys_section , self . axis ) <EOL> @ attr . gpu <EOL> def test_forward_gpu ( self ) : <EOL> self . check_forward ( <EOL> cuda . to_gpu ( self . x ) , <EOL> [ cuda . to_gpu ( y . copy ( ) ) for y in self . ys ] , <EOL> self . ys_section , axis = self . axis ) <EOL> def check_backward ( self , x_data , indices_or_sections , axis ) : <EOL> x = chainer . Variable ( x_data ) <EOL> ys = functions . split_axis ( x , indices_or_sections , axis ) <EOL> for y in ys : <EOL> y . grad = y . data <EOL> ys [ <NUM_LIT:0> ] . backward ( ) <EOL> gradient_check . assert_allclose ( x . data , x . grad , atol = <NUM_LIT:0> , rtol = <NUM_LIT:0> ) <EOL> def test_backward_cpu ( self ) : <EOL> self . check_backward ( self . x , self . ys_section , axis = self . axis ) <EOL> @ attr . gpu <EOL> def test_backward_gpu ( self ) : <EOL> self . check_backward ( <EOL> cuda . to_gpu ( self . x ) , self . ys_section , axis = self . axis ) <EOL> class TestSplitAxisNone ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . x = numpy . array ( [ <NUM_LIT:1> , <NUM_LIT:2> ] , dtype = numpy . float32 ) <EOL> self . ys_section = [ <NUM_LIT:1> ] <EOL> self . axis = <NUM_LIT:0> <EOL> def check_backward ( self , x_data , indices_or_sections , axis ) : <EOL> x = chainer . Variable ( x_data ) <EOL> ys = functions . split_axis ( x , indices_or_sections , axis ) <EOL> ys [ <NUM_LIT:0> ] . grad = ys [ <NUM_LIT:0> ] . data <EOL> ys [ <NUM_LIT:0> ] . backward ( ) <EOL> gx = numpy . array ( [ <NUM_LIT:1> , <NUM_LIT:0> ] ) <EOL> gradient_check . assert_allclose ( gx , x . grad , atol = <NUM_LIT:0> , rtol = <NUM_LIT:0> ) <EOL> def test_backward_cpu ( self ) : <EOL> self . check_backward ( self . x , self . ys_section , axis = self . axis ) <EOL> @ attr . gpu <EOL> def test_backward_gpu ( self ) : <EOL> self . check_backward ( <EOL> cuda . to_gpu ( self . x ) , self . ys_section , axis = self . axis ) <EOL> class TestSplitAxisForceArray ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> self . x = numpy . arange ( <NUM_LIT> , dtype = numpy . float32 ) . reshape ( <NUM_LIT:2> , <NUM_LIT:7> , <NUM_LIT:3> ) <EOL> self . axis = <NUM_LIT:1> <EOL> def check_forward_force_tuple ( self , x_data , axis ) : <EOL> x = chainer . Variable ( x_data ) <EOL> ys = functions . split_axis ( x , <NUM_LIT:1> , axis , force_tuple = True ) <EOL> self . assertIsInstance ( ys , tuple ) <EOL> self . assertEqual ( len ( ys ) , <NUM_LIT:1> ) <EOL> def test_forward_force_tuple_cpu ( self ) : <EOL> self . check_forward_force_tuple ( self . x , self . axis ) <EOL> @ attr . gpu <EOL> def test_forward_force_tuple_gpu ( self ) : <EOL> self . check_forward_force_tuple ( cuda . to_gpu ( self . x ) , axis = self . axis ) <EOL> def check_forward_single ( self , x_data , axis ) : <EOL> x = chainer . Variable ( x_data ) <EOL> ys = functions . split_axis ( x , <NUM_LIT:1> , axis ) <EOL> self . assertIsInstance ( ys , chainer . Variable ) <EOL> def test_forward_single_cpu ( self ) : <EOL> self . check_forward_single ( self . x , self . axis ) <EOL> @ attr . gpu <EOL> def test_forward_single_gpu ( self ) : <EOL> self . check_forward_single ( cuda . to_gpu ( self . x ) , axis = self . axis ) <EOL> testing . run_module ( __name__ , __file__ ) </s>
